#任务执行者
from app_utils import websearch
from sshcode.plan_execute.model_utils import getLLM

_system_prompt_template = '''
根据任务调用工具查询用户信息
不需要其他信息
工具
{tools_name}
'''

_human_prompt_template='''
当前子任务:
{task_name}
'''

from langchain_core.prompts import ChatPromptTemplate
from langgraph.prebuilt import create_react_agent
from langchain_core.output_parsers import StrOutputParser
from langchain_core.tools import tool
from typing import Annotated

@tool
def web_searcher(query:Annotated[str,'需要查询的子任务']):
    '''调用互联网工具，查询信息'''
    return websearch(query)

class Executor:
    def __init__(self,llm):
        _tools = [web_searcher]
        _prompt = ChatPromptTemplate.from_messages([
            ("system",_system_prompt_template),
            ("human",_human_prompt_template)
        ])
        # _prompt = _prompt.partial(tools_name=",".join([_tool.name for _tool in _tools]))
        _prompt = _prompt.partial(tools_name=", ".join([_tool.name for _tool in _tools]))
        agent = create_react_agent(llm, _tools)
        self._parser = StrOutputParser()
        self._chain = _prompt | agent

    def __call__(self, state):
        _rt = self._chain.invoke(state)

        return self._parser.invoke(_rt['messages'][-1])

if __name__ == '__main__':
    llm = getLLM()
    _exec = Executor(llm)
    _rt = _exec({"task_name":"查找2024年法国奥运会10米跳水项目的冠军"})
    print(_rt)